import requests
import json
import re
import base64
from datetime import datetime, timedelta


API_KEY = "sk-*"  
API_URL = "https://api.deepseek.com/v1/chat/completions"

HEADERS = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {API_KEY}"
}

# helper function 1: encode image to base64
def encode_image_to_base64(filepath):
    try:
        with open(filepath, "rb") as f:
            encoded = base64.b64encode(f.read()).decode("utf-8")
        return encoded
    except FileNotFoundError:
        print(f"[错误] 图片路径不存在: {filepath}")
        return ""

# helper function 2: extract json from response
def extract_json_from_response(response_text):
    match = re.search(r"```json\s*(.*?)\s*```", response_text, re.DOTALL)
    if match:
        json_str = match.group(1)
    else:
        json_str = response_text
    try:
        return json.loads(json_str)
    except json.JSONDecodeError as e:
        print("无法解析 JSON：", e)
        print("原始内容：\n", response_text)
        return None
    

def classify_events(events):
    CLASSIFY_PROMPT = """
你是一个安全巡检助手，请根据每个事件的描述（summary）和内容（messages），判断该事件属于以下哪一类：

- 巡检问题：发现现场隐患或风险点
- 验收：施工完成后的检查
- 旁站：施工过程中的现场陪同监督
- 闭环：对之前问题的处理、确认
- 其他：不属于以上类型的其他事项

请使用如下 JSON 格式返回每个事件的分类：

[
  {
    "event_id": 1,
    "class": "巡检问题"
  },
  ...
]
"""

    event_summaries = [
        {
            "event_id": e["event_id"],
            "summary": e["summary"],
            "messages": e["messages"]
        } for e in events
    ]

    prompt = CLASSIFY_PROMPT + "\n以下是事件信息：\n" + json.dumps(event_summaries, ensure_ascii=False, indent=2)

    payload = {
        "model": "deepseek-chat",
        "messages": [
            {"role": "system", "content": "你是一个分类助手"},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.2
    }

    try:
        response = requests.post(API_URL, headers=HEADERS, json=payload)
        if response.status_code == 200:
            return extract_json_from_response(response.json()["choices"][0]["message"]["content"])
        else:
            print("分类请求失败:", response.status_code, response.text)
            return None
    except Exception as e:
        print("分类 API 调用失败:", str(e))
        return None


def add_class_to_events(events, classifications):
    classification_dict = {c["event_id"]: c["class"] for c in classifications}
    for event in events:
        event["class"] = classification_dict.get(event["event_id"], "其他")
    return events

def main():
    img2_base64 = encode_image_to_base64("/Users/wangzijian/Desktop/实习gauz/gauz_ai_clustering/IMG_2.jpg")
    img3_base64 = encode_image_to_base64("/Users/wangzijian/Desktop/实习gauz/gauz_ai_clustering/image03.jpg")

    events = [
        {
            "event_id": 1,
            "summary": "在3号楼东南角发现脚手架松动",
            "messages": [
                {
                    "type": "text",
                    "content": "在3号楼东南角发现脚手架松动，有坠落风险",
                    "timestamp": "2024-04-01T08:30:00Z"
                },
                {
                    "type": "text",
                    "content": "责任人张三已通知整改",
                    "timestamp": "2024-04-01T08:31:10Z"
                }
            ],
            "candidate_images": []
        },
        {
            "event_id": 2,
            "summary": "2号楼发现电缆裸露",
            "messages": [
                {
                    "type": "text",
                    "content": "2号楼发现电缆裸露，存在触电危险",
                    "timestamp": "2024-04-01T10:00:00Z"
                }
            ],
            "candidate_images": [
                {
                    "base64": img2_base64,
                    "filename": "IMG_2.jpg",
                    "timestamp": "2024-04-01T10:01:00Z"
                }
            ]
        },

        {
            "event_id": 3,
            "summary": "5号楼地下室防水层已完成，准备验收",
            "messages": [
            {
                "type": "text",
                "content": "5号楼地下室防水层施工于昨日完成，今日进行初步验收",
                "timestamp": "2024-04-02T08:45:00Z"
            },
            {
                "type": "text",
                "content": "验收人员已到场，准备进行验收",
                "timestamp": "2024-04-02T08:50:00Z"
            }
        ],
            "candidate_images": []
        },

        {
            "event_id": 4,
            "summary": "钢筋绑扎作业进行中，监理确认规范",
            "messages": [
            {
                "type": "text",
                "content": "今日上午9点，3号楼主体钢筋绑扎作业开始",
                "timestamp": "2024-04-01T09:00:00Z"
            },
            {
                "type": "text",
                "content": "监理张工已到场确认施工规范",
                "timestamp": "2024-04-01T09:15:00Z"
            }
        ],
            "candidate_images": []
        }
    ]

    # 分类调用
    classifications = classify_events(events)
    if classifications is None:
        print("分类失败，未返回结果")
        return

    events_with_class = add_class_to_events(events, classifications)

    # 打印结果
    print("\n分类结果：")
    for e in events_with_class:
        print(f"- Event {e['event_id']}: 类型 = {e['class']}，摘要 = {e['summary']}")


# 调用主函数
if __name__ == "__main__":
    main()
